#!/usr/local/bin/python
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import IUPAC
from multiprocessing import cpu_count
import gsat,os,sys,re,hmmtop,pickle,os,shutil,scalagsat,threading,scala
#os.environ['BIOV_DEBUG']='True'
class build():
	
	def __init__(self):
		self.split = 6
		self.fasta_file = False
		self.flank = 10
		self.tms = []
		self.loops ={}
		self.outfile = "horizontal.txt"
		self.outdir = False
		self.min_tms = 10
		self.threads = cpu_count()
		self.method = 3
		self.VSrestrict = False
		self.VTrestrict = False
		self.consecutive = False
		self.fasta_only = False
		self.cc = False
		
	def __call__(self):
		self.hmmtop()
		self.loops=self.build_flanks()
		try:
			os.mkdir("./"+self.outdir)
		except:
			pass
		'''	
		if self.cc is True:
			pass
			self.clustal_correct()
			quit()
		'''
		self.pool = scala.ThreadPool(self.threads)
		if self.method==1 or self.method==3:
			print "# Searching for horizontal TMS repeats"
			self.pool.queueTask(self.intercompare,None)
		if self.method==2 or self.method==3:
			print "# Searching for Vertical TMS repeats"
			self.vertical()
		self.pool.joinAll()
		print "...Running in-depth repeat search on %i Cores\n" %(self.threads)
		print "You may view results in realtime by CD'ing to your output directory"
		print "And typing readlive.py <results_file>\n"
		print "You may also view any complete alignment by typing: show_alignment.py"
		print "When prompted, paste any ONE line from your results file."
		print "Be sure to preform this in a new terminal window, to prevent interrupting this current scan."
		
	def hmmtop(self):
		hm = hmmtop.tools()
		hm.add_library("SF", self.fasta_file.name)
		hm.scan_libraries()
		self.tms = hm.results['SF']
		return
		
	def get_groups(self,keys):
		db = {}
		for size in keys:
			for row in keys:
				end = row+(size-1)
				if end > len(keys):
					break
				db.setdefault(size,[]).append(range(row,end+1))
		return db
		
	def build_flanks(self):
		self.fasta_file.seek(0)
		fastas = SeqIO.to_dict(SeqIO.parse(self.fasta_file,'fasta'))
		myloops = {}
		for symbol, tmss in self.tms.items():
			in_use=[]
			[in_use.extend(range(i[0],i[1]+1)) for i in tmss.values()]
			for tms, residues in tmss.items():
				maxlen = len(str(fastas[symbol].seq))
				for extend in range(self.flank):
					left = residues[0] - 1
					right = residues[1] + 1
					# check validity
					if left >= 1 and left not in in_use:
						in_use.append(left)
						residues[0] = left
					if right <= maxlen and right not in in_use:
						in_use.append(right)
						residues[1] = right
				loop = fastas[symbol].seq[residues[0]-1:residues[1]]
				myloops.setdefault(symbol,[]).append(loop)
		return myloops
	'''			
	def clustal_correct():
		flank = copy(self.flank)
		self.flank = 0
		loops = self.build_flanks()
		self.flank = flank
		#align = open("alignment.faa",'wb+')
		for loop in loops:
			print loop
			#record = SeqRecord(Seq(seq,IUPAC.protien),id=)
	'''

	def intercompare(self,td=None): # Preforms a LINEAR search
		# Extract TMS requirements
		if self.fasta_only is True:
			return
		results = []
		specific_loops={}
		tscore = 0
		myloops = [i for i in self.loops.items()]
		handle = open(self.outdir+"/"+self.outfile,'wb+')
		score = gsat.cmd()
		for symbol, tmss in myloops:
			keys = range(1,len(tmss)+1)
			try:
				groups = self.get_groups(keys)[self.split]
			except:
				continue
			for i,tms in enumerate(groups):
				subject = tms
				#target = groups[i+self.split] if (i+self.split) <= len(groups)-1 else False
				#if target is False:
				#	continue
				for target in groups[i+self.split:]:
					#'''
					try:
						subject_seq = ''.join([str(self.loops[symbol][int(i)-1]) for i in subject])
						target_seq = ''.join([str(self.loops[symbol][int(i)-1]) for i in target])
						#score = gsat.compare(subject_seq,target_seq)
						score.asequence=subject_seq
						score.bsequence=target_seq
						score()
						subject = [str(i) for i in subject]
						gaps = score.gaps
						target = [str(i) for i in target]
						zones = '-'.join(subject)+' && '+'-'.join(target)
						line = "%s\t%s\t%s\t%s" %(symbol,zones,str(gaps),str(round(int(score.zscore))))
						handle.write(line+"\n")
						handle.flush()
						results.append((symbol,zones,int(gaps),round(int(score.zscore))))
					except:
						pass
					#'''
		handle = open(self.outdir+"/"+self.outfile,'wb+')
		results.sort(key=lambda x:(100-float(x[-2]),float(x[-1])),reverse=True)
		for i in results:
			row = "\t".join(str(x) for x in i)
			handle.write(row+"\n")
		handle.flush()
		return self.outfile
		
	def vertical(self):
		myloops = [i for i in self.loops.items() if len(i[1]) >= self.min_tms and len(i[1]) <=self.max_tms ]
		mydirs =[]
		for symbol, tms in myloops:
			keys = range(1,len(tms)+1)
			groups = self.get_groups(keys)[self.split]
			for x,tms in enumerate(groups):
				previous = []
				subject = tms
				if self.VSrestrict is not False and tms not in self.VSrestrict:
					continue
				subjectdir = "./"+self.outdir+"/"+"-".join([str(i) for i in subject])
				mydirs.append(subjectdir)
				try:
					os.makedirs(subjectdir)
				except:
					pass
				# Write Subject FASTA to file
				subject_file = open(subjectdir+"/subject.faa",'ab+')
				subject_seq = ''.join([str(self.loops[symbol][int(i)-1]) for i in subject])
				subject_desc = 'TMS #(%s)'%(",".join([str(i) for i in subject]))
				record = SeqRecord(Seq(subject_seq,IUPAC.protein),id=symbol,description=subject_desc)
				SeqIO.write(record,subject_file,'fasta')
				subject_file.flush()
				# Iterate through Target candidates and write them as we go
				target_file = open(subjectdir+'/target.faa','ab+')
				for targets in groups[x+self.split:]:
					if bool(set(targets) & set(previous)) and self.consecutive is True:
						continue
					if self.VTrestrict is not False and targets not in self.VTrestrict:
						continue
					target_seq = ''.join([str(self.loops[symbol][int(i)-1]) for i in targets])
					target_desc = 'TMS #(%s)'%(",".join([str(i) for i in targets]))
					record = SeqRecord(Seq(target_seq,IUPAC.protein),id=symbol,description=target_desc)
					SeqIO.write(record,target_file,'fasta')
					previous = targets
				target_file.flush()
		# Delete dirs with empty
		mydirs = list(set(mydirs))
		good_dirs = []
		if self.fasta_only is True:
			return
		for directory in mydirs:
			if os.path.getsize(directory+"/target.faa") is 0:
				shutil.rmtree(directory)
			else:
				good_dirs.append(directory)
				
		# Run Scala_GSAT on all Good dirs.
		print "# Scala-GSAT is in progress...Check vertical.txt:"
		for thisdir in good_dirs:
			#threading.Thread(target=self.sub_scala,args=(thisdir+'/subject.faa',thisdir+'/target.faa',thisdir+'/results.txt',)).start()
			self.pool.queueTask(self.sub_scala,(thisdir+'/subject.faa',thisdir+'/target.faa',thisdir+'/vertical.txt'))
		return
		
			
			
	def sub_scala(self,data):
		(subject,target,outfile) = data
		scala = scalagsat.compare()
		scala.subject = subject
		scala.target = target
		scala.outfile = outfile
		scala.quiet = True
		scala.defline = True
		scala()

		
if __name__=='__main__':
	from optparse import OptionParser,OptionGroup
	desc = "This is the Ancient program. This tool will find very old (and new) TMS repeat units. Results can be viewed as they become\
	available by using the readlive.py command. Alignments can be viewed using the show_alignment.py command. This tool was created by:\
	Vamsee Reddy, 2010 Saierlab <Symphony.Dev@gmail.com> :: http://projectv.tcdb.org"
	version = "Ancient Repeat Tool v1.0"
	opts = OptionParser(description=desc,version=version)
	opts.add_option('-i',action='store',type='string',dest='input', help="Path to your subject fasta file")
	opts.add_option('-r',action='store',type='int',dest='repeat', help="Repeat Size")
	opts.add_option('-o',action='store',type='string',dest='output', help="Output Name")
	optional = OptionGroup(opts,"Optional Settings")
	optional.add_option('--min',action='store',type='int',dest='min', default=0,help="Minimum TMS Requirement")
	optional.add_option('--max',action='store',type='int',dest='max', default=500,help="Maximum TMS Requirement")
	optional.add_option('--flank',action='store',type='int',dest='f_length', default=10,help="Size of hydrophillic padding around TMSs (10 aa)")
	optional.add_option('--method',action='store',type='int',dest='method',default=3,help="[1] horizontal, [2] Vertical, [3] Both ")
	optional.add_option('--VSrestrict',action='store',type='string',dest='sset',default=False,help="Vertical Subject Restriction.")
	optional.add_option('--VTrestrict',action='store',type='string',dest='tset',default=False,help="Vertical Target Restriction.")
	optional.add_option('--consecutive',action='store_true',default=False,help="Includes only consecutive targets 'r' TMSs apart")
	optional.add_option('--clustalcorrect',action='store_true',default=False,help="Correct different numbers of TMSs with a multiple alignment")
	optional.add_option('--fasta_only',action='store_true',default=False, help="Build ONLY Fasta DB(s) (No TMS Searching)")
	optional.add_option('--threads',action='store',type='int',dest='threads',default=cpu_count(), help="Number of threads <default:%i>"%cpu_count())
	
	opts.add_option_group(optional)				
	(cli,args)=opts.parse_args()
	r = build()
	try:
		
		r.split=cli.repeat
		r.min_tms=cli.min
		r.max_tms=cli.max
		r.fasta_file=open(cli.input,'r')
		r.threads=cli.threads
		r.method=cli.method
		r.outdir=cli.output
		r.flank = cli.f_length
		r.fasta_only = cli.fasta_only
		r.consecutive = cli.consecutive
		r.cc = cli.clustalcorrect
		restrict = re.compile('\(([0-9, ]+)\)')
		if cli.sset is False:
			r.VSrestrict = False
		else:
			restricted=restrict.findall(cli.sset)
			restricted = [[int(x.strip()) for x in i.split(',')] for i in restricted]
			r.VSrestrict = restricted
		if cli.tset is False:
			r.VTrestrict = False
		else:
			restricted=restrict.findall(cli.tset)
			restricted = [[int(x.strip()) for x in i.split(',')] for i in restricted]
			r.VTrestrict = restricted
		r()
	except:
		opts.print_help()